{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:35.553282Z",
     "start_time": "2024-03-06T10:37:35.548266Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from factorVAE.factor_VAE import FactorVAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:35.585328Z",
     "start_time": "2024-03-06T10:37:35.574285Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "device = \"cuda\"\n",
    "stock = \"sp500\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:36.906382Z",
     "start_time": "2024-03-06T10:37:35.644103Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total step: 948\n",
      "Time span: 60\n",
      "Stock size: 487\n",
      "Feature size: 5\n"
     ]
    }
   ],
   "source": [
    "features_dataset = torch.load(f\"./dataset/{stock}/train/feat.pt\")\n",
    "returns_dataset = torch.load(f\"./dataset/{stock}/train/ret.pt\")\n",
    "val_features_dataset = torch.load(f\"./dataset/{stock}/val/feat.pt\")\n",
    "val_returns_dataset = torch.load(f\"./dataset/{stock}/val/ret.pt\")\n",
    "test_features_dataset = torch.load(f\"./dataset/{stock}/test/feat.pt\")\n",
    "test_returns_dataset = torch.load(f\"./dataset/{stock}/test/ret.pt\")\n",
    "\n",
    "print(f\"Total step: {features_dataset.shape[0]}\")\n",
    "print(f\"Time span: {features_dataset.shape[1]}\")\n",
    "print(f\"Stock size: {features_dataset.shape[2]}\")\n",
    "print(f\"Feature size: {features_dataset.shape[3]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:36.921921Z",
     "start_time": "2024-03-06T10:37:36.909383Z"
    }
   },
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "characteristic_size = features_dataset.shape[3]\n",
    "stock_size = features_dataset.shape[2]\n",
    "latent_size = 64\n",
    "factor_size = 32\n",
    "time_span = features_dataset.shape[1]\n",
    "gru_input_size = 64\n",
    "hidden_size = 64\n",
    "lr = 1e-5\n",
    "epochs = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:36.938019Z",
     "start_time": "2024-03-06T10:37:36.924977Z"
    }
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "def get_dataloader(data, label, device=device, batch_size=batch_size):\n",
    "    data = torch.Tensor(data).to(device)\n",
    "    label = torch.Tensor(label).to(device).long()\n",
    "    ds = TensorDataset(data, label)\n",
    "    dl = DataLoader(ds, batch_size=batch_size, shuffle=True)\n",
    "    return dl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:37.282404Z",
     "start_time": "2024-03-06T10:37:36.942054Z"
    }
   },
   "outputs": [],
   "source": [
    "train_dl = get_dataloader(features_dataset, returns_dataset)\n",
    "val_dl = get_dataloader(val_features_dataset, val_returns_dataset)\n",
    "test_dl = get_dataloader(test_features_dataset, test_returns_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:37.343707Z",
     "start_time": "2024-03-06T10:37:37.283455Z"
    }
   },
   "outputs": [],
   "source": [
    "factor_VAE = FactorVAE(\n",
    "    characteristic_size=characteristic_size,\n",
    "    stock_size=stock_size,\n",
    "    latent_size=latent_size,\n",
    "    factor_size=factor_size,\n",
    "    time_span=time_span,\n",
    "    gru_input_size=gru_input_size,\n",
    "    hidden_size=hidden_size\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:37.359711Z",
     "start_time": "2024-03-06T10:37:37.345708Z"
    }
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(factor_VAE.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:37.374982Z",
     "start_time": "2024-03-06T10:37:37.361885Z"
    }
   },
   "outputs": [],
   "source": [
    "def train_loop(dataloader, model, optimizer):\n",
    "    for batch, (feat, ret) in enumerate(dataloader):\n",
    "        loss = model.run_model(feat, ret, gamma=1)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if (batch + 1) % 20 == 0:\n",
    "            print(f\"batch: {batch}, loss: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def val_loop(dataloader, model):\n",
    "    total_loss = 0\n",
    "    total_batches = 0\n",
    "    with torch.no_grad():  # 禁用梯度计算\n",
    "        for batch, (feat, ret) in enumerate(dataloader):\n",
    "            loss = model.run_model(feat, ret, gamma=1)\n",
    "\n",
    "            total_loss += loss.item()\n",
    "            total_batches += 1\n",
    "            \n",
    "    average_loss = total_loss / total_batches\n",
    "    print(f\"Validation Loss: {average_loss}\")\n",
    "    return average_loss"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T10:37:37.390562Z",
     "start_time": "2024-03-06T10:37:37.377067Z"
    }
   },
   "execution_count": 21
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:22.102556Z",
     "start_time": "2024-03-06T10:37:37.392651Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Epoch: 0 ===\n",
      "batch: 19, loss: 1.970694899559021\n",
      "batch: 39, loss: 1.6346464157104492\n",
      "batch: 59, loss: 0.8889301419258118\n",
      "Validation Loss: 1.3166942993799846\n",
      "=== Epoch: 1 ===\n",
      "batch: 19, loss: 1.258047342300415\n",
      "batch: 39, loss: 1.0256068706512451\n",
      "batch: 59, loss: 0.7358337044715881\n",
      "Validation Loss: 0.873470253414578\n",
      "=== Epoch: 2 ===\n",
      "batch: 19, loss: 0.8524231314659119\n",
      "batch: 39, loss: 0.8250256776809692\n",
      "batch: 59, loss: 0.6811049580574036\n",
      "Validation Loss: 0.7880741424030728\n",
      "=== Epoch: 3 ===\n",
      "batch: 19, loss: 0.7838475108146667\n",
      "batch: 39, loss: 0.7645266056060791\n",
      "batch: 59, loss: 0.6326335668563843\n",
      "Validation Loss: 0.7307743430137634\n",
      "=== Epoch: 4 ===\n",
      "batch: 19, loss: 0.727276086807251\n",
      "batch: 39, loss: 0.7098177671432495\n",
      "batch: 59, loss: 0.5850622653961182\n",
      "Validation Loss: 0.6784059007962545\n",
      "=== Epoch: 5 ===\n",
      "batch: 19, loss: 0.6367621421813965\n",
      "batch: 39, loss: 0.6105685830116272\n",
      "batch: 59, loss: 0.525255560874939\n",
      "Validation Loss: 0.583411349190606\n",
      "=== Epoch: 6 ===\n",
      "batch: 19, loss: 0.5732369422912598\n",
      "batch: 39, loss: 0.5532146692276001\n",
      "batch: 59, loss: 0.47177594900131226\n",
      "Validation Loss: 0.5260400109820895\n",
      "=== Epoch: 7 ===\n",
      "batch: 19, loss: 0.4754277765750885\n",
      "batch: 39, loss: 0.4495405852794647\n",
      "batch: 59, loss: 0.404153048992157\n",
      "Validation Loss: 0.4254729714658525\n",
      "=== Epoch: 8 ===\n",
      "batch: 19, loss: 0.4074650704860687\n",
      "batch: 39, loss: 0.3847179412841797\n",
      "batch: 59, loss: 0.33932897448539734\n",
      "Validation Loss: 0.3580361538463169\n",
      "=== Epoch: 9 ===\n",
      "batch: 19, loss: 0.3382859230041504\n",
      "batch: 39, loss: 0.3115924298763275\n",
      "batch: 59, loss: 0.2644653022289276\n",
      "Validation Loss: 0.28459302915467155\n",
      "=== Epoch: 10 ===\n",
      "batch: 19, loss: 0.2586769461631775\n",
      "batch: 39, loss: 0.22856417298316956\n",
      "batch: 59, loss: 0.17776153981685638\n",
      "Validation Loss: 0.19482573370138803\n",
      "=== Epoch: 11 ===\n",
      "batch: 19, loss: 0.16361117362976074\n",
      "batch: 39, loss: 0.1296946108341217\n",
      "batch: 59, loss: 0.07595633715391159\n",
      "Validation Loss: 0.09087514463398191\n",
      "=== Epoch: 12 ===\n",
      "batch: 19, loss: 0.05662503093481064\n",
      "batch: 39, loss: 0.017510954290628433\n",
      "batch: 59, loss: -0.03802074119448662\n",
      "Validation Loss: -0.02382000494334433\n",
      "=== Epoch: 13 ===\n",
      "batch: 19, loss: -0.06387991458177567\n",
      "batch: 39, loss: -0.10738597810268402\n",
      "batch: 59, loss: -0.16810822486877441\n",
      "Validation Loss: -0.15674827496210733\n",
      "=== Epoch: 14 ===\n",
      "batch: 19, loss: -0.20184321701526642\n",
      "batch: 39, loss: -0.2516070604324341\n",
      "batch: 59, loss: -0.316500723361969\n",
      "Validation Loss: -0.3064582082960341\n",
      "=== Epoch: 15 ===\n",
      "batch: 19, loss: -0.3563210070133209\n",
      "batch: 39, loss: -0.4116831123828888\n",
      "batch: 59, loss: -0.48165127635002136\n",
      "Validation Loss: -0.47238027056058246\n",
      "=== Epoch: 16 ===\n",
      "batch: 19, loss: -0.5287147760391235\n",
      "batch: 39, loss: -0.5926057696342468\n",
      "batch: 59, loss: -0.6709787845611572\n",
      "Validation Loss: -0.6643386946784126\n",
      "=== Epoch: 17 ===\n",
      "batch: 19, loss: -0.7309073209762573\n",
      "batch: 39, loss: -0.8044073581695557\n",
      "batch: 59, loss: -0.890590488910675\n",
      "Validation Loss: -0.8841789960861206\n",
      "=== Epoch: 18 ===\n",
      "batch: 19, loss: -0.9601172804832458\n",
      "batch: 39, loss: -1.0450541973114014\n",
      "batch: 59, loss: -1.1426643133163452\n",
      "Validation Loss: -1.1383753220240276\n",
      "=== Epoch: 19 ===\n",
      "batch: 19, loss: -1.2269713878631592\n",
      "batch: 39, loss: -1.320830225944519\n",
      "batch: 59, loss: -1.4299310445785522\n",
      "Validation Loss: -1.4264234171973333\n",
      "=== Epoch: 20 ===\n",
      "batch: 19, loss: -1.524030089378357\n",
      "batch: 39, loss: -1.6305716037750244\n",
      "batch: 59, loss: -1.7481422424316406\n",
      "Validation Loss: -1.7476573917600844\n",
      "=== Epoch: 21 ===\n",
      "batch: 19, loss: -1.8585929870605469\n",
      "batch: 39, loss: -1.9802772998809814\n",
      "batch: 59, loss: -2.1131575107574463\n",
      "Validation Loss: -2.1098955737219915\n",
      "=== Epoch: 22 ===\n",
      "batch: 19, loss: -2.2342703342437744\n",
      "batch: 39, loss: -2.3677148818969727\n",
      "batch: 59, loss: -2.5113418102264404\n",
      "Validation Loss: -2.51106686062283\n",
      "=== Epoch: 23 ===\n",
      "batch: 19, loss: -2.6456894874572754\n",
      "batch: 39, loss: -2.7905876636505127\n",
      "batch: 59, loss: -2.942988634109497\n",
      "Validation Loss: -2.9452608956231012\n",
      "=== Epoch: 24 ===\n",
      "batch: 19, loss: -3.091737747192383\n",
      "batch: 39, loss: -3.2471957206726074\n",
      "batch: 59, loss: -3.411412000656128\n",
      "Validation Loss: -3.415555953979492\n",
      "=== Epoch: 25 ===\n",
      "batch: 19, loss: -3.573369264602661\n",
      "batch: 39, loss: -3.743648052215576\n",
      "batch: 59, loss: -3.9213926792144775\n",
      "Validation Loss: -3.927472644382053\n",
      "=== Epoch: 26 ===\n",
      "batch: 19, loss: -4.096986293792725\n",
      "batch: 39, loss: -4.280220985412598\n",
      "batch: 59, loss: -4.472455978393555\n",
      "Validation Loss: -4.4745713869730634\n",
      "=== Epoch: 27 ===\n",
      "batch: 19, loss: -4.6596221923828125\n",
      "batch: 39, loss: -4.83982515335083\n",
      "batch: 59, loss: -5.043643474578857\n",
      "Validation Loss: -5.04758686489529\n",
      "=== Epoch: 28 ===\n",
      "batch: 19, loss: -5.233639717102051\n",
      "batch: 39, loss: -5.434157371520996\n",
      "batch: 59, loss: -5.611378192901611\n",
      "Validation Loss: -5.617280377282037\n",
      "=== Epoch: 29 ===\n",
      "batch: 19, loss: -5.761894226074219\n",
      "batch: 39, loss: -5.924764156341553\n",
      "batch: 59, loss: -6.143333435058594\n",
      "Validation Loss: -6.145766099294026\n",
      "=== Epoch: 30 ===\n",
      "batch: 19, loss: -6.241183757781982\n",
      "batch: 39, loss: -6.44115686416626\n",
      "batch: 59, loss: -6.431252479553223\n",
      "Validation Loss: -6.415894190470378\n",
      "=== Epoch: 31 ===\n",
      "batch: 19, loss: -6.349624156951904\n",
      "batch: 39, loss: -6.523405075073242\n",
      "batch: 59, loss: -6.694562911987305\n",
      "Validation Loss: -6.704693582322863\n",
      "=== Epoch: 32 ===\n",
      "batch: 19, loss: -6.850001335144043\n",
      "batch: 39, loss: -7.007874011993408\n",
      "batch: 59, loss: -6.865318298339844\n",
      "Validation Loss: -6.700658957163493\n",
      "=== Epoch: 33 ===\n",
      "batch: 19, loss: -6.961499214172363\n",
      "batch: 39, loss: -7.206525802612305\n",
      "batch: 59, loss: -6.726613521575928\n",
      "Validation Loss: -6.146325959099664\n",
      "=== Epoch: 34 ===\n",
      "batch: 19, loss: -6.421561241149902\n",
      "batch: 39, loss: -6.682771682739258\n",
      "batch: 59, loss: -6.997725486755371\n",
      "Validation Loss: -7.0122241444057885\n",
      "=== Epoch: 35 ===\n",
      "batch: 19, loss: -7.204143047332764\n",
      "batch: 39, loss: -7.383072853088379\n",
      "batch: 59, loss: -7.481448650360107\n",
      "Validation Loss: -7.488672521379259\n",
      "=== Epoch: 36 ===\n",
      "batch: 19, loss: -7.497856616973877\n",
      "batch: 39, loss: -7.693999767303467\n",
      "batch: 59, loss: -6.71256160736084\n",
      "Validation Loss: -6.678402900695801\n",
      "=== Epoch: 37 ===\n",
      "batch: 19, loss: -6.499363899230957\n",
      "batch: 39, loss: -6.131419658660889\n",
      "batch: 59, loss: -6.383765697479248\n",
      "Validation Loss: -6.423461861080593\n",
      "=== Epoch: 38 ===\n",
      "batch: 19, loss: -6.879788398742676\n",
      "batch: 39, loss: -7.0786519050598145\n",
      "batch: 59, loss: -7.238460540771484\n",
      "Validation Loss: -7.281897226969401\n",
      "=== Epoch: 39 ===\n",
      "batch: 19, loss: -7.474748611450195\n",
      "batch: 39, loss: -7.668832302093506\n",
      "batch: 59, loss: -7.867793560028076\n",
      "Validation Loss: -7.875810994042291\n",
      "=== Epoch: 40 ===\n",
      "batch: 19, loss: -7.983368873596191\n",
      "batch: 39, loss: -7.918956756591797\n",
      "batch: 59, loss: -7.975287437438965\n",
      "Validation Loss: -7.9812683529324\n",
      "=== Epoch: 41 ===\n",
      "batch: 19, loss: -7.942657470703125\n",
      "batch: 39, loss: -8.142461776733398\n",
      "batch: 59, loss: -5.739562511444092\n",
      "Validation Loss: -5.189321253034803\n",
      "=== Epoch: 42 ===\n",
      "batch: 19, loss: -3.9593212604522705\n",
      "batch: 39, loss: -4.195781707763672\n",
      "batch: 59, loss: -4.910943031311035\n",
      "Validation Loss: -4.98177613152398\n",
      "=== Epoch: 43 ===\n",
      "batch: 19, loss: -4.792354106903076\n",
      "batch: 39, loss: -4.647397994995117\n",
      "batch: 59, loss: -4.037032127380371\n",
      "Validation Loss: -4.038232856326633\n",
      "=== Epoch: 44 ===\n",
      "batch: 19, loss: -4.208652496337891\n",
      "batch: 39, loss: -4.600215911865234\n",
      "batch: 59, loss: -6.262112140655518\n",
      "Validation Loss: -6.370209852854411\n",
      "=== Epoch: 45 ===\n",
      "batch: 19, loss: -4.868869304656982\n",
      "batch: 39, loss: -6.433544635772705\n",
      "batch: 59, loss: -6.4619340896606445\n",
      "Validation Loss: -6.504967901441786\n",
      "=== Epoch: 46 ===\n",
      "batch: 19, loss: -6.572746753692627\n",
      "batch: 39, loss: -6.648965358734131\n",
      "batch: 59, loss: -6.741440773010254\n",
      "Validation Loss: -6.749618159400092\n",
      "=== Epoch: 47 ===\n",
      "batch: 19, loss: -6.847076892852783\n",
      "batch: 39, loss: -6.982706069946289\n",
      "batch: 59, loss: -7.159943103790283\n",
      "Validation Loss: -7.17451540629069\n",
      "=== Epoch: 48 ===\n",
      "batch: 19, loss: -7.410120487213135\n",
      "batch: 39, loss: -7.7335991859436035\n",
      "batch: 59, loss: -7.9977707862854\n",
      "Validation Loss: -8.014525201585558\n",
      "=== Epoch: 49 ===\n",
      "batch: 19, loss: -8.205436706542969\n",
      "batch: 39, loss: -7.751246929168701\n",
      "batch: 59, loss: -8.225507736206055\n",
      "Validation Loss: -8.244175487094456\n",
      "=== Epoch: 50 ===\n",
      "batch: 19, loss: -6.190258026123047\n",
      "batch: 39, loss: -4.611129283905029\n",
      "batch: 59, loss: -6.097648620605469\n",
      "Validation Loss: -6.433694362640381\n",
      "=== Epoch: 51 ===\n",
      "batch: 19, loss: -6.561679840087891\n",
      "batch: 39, loss: -6.605219841003418\n",
      "batch: 59, loss: -6.722774028778076\n",
      "Validation Loss: -6.728431807623969\n",
      "=== Epoch: 52 ===\n",
      "batch: 19, loss: -6.771795749664307\n",
      "batch: 39, loss: -6.844500541687012\n",
      "batch: 59, loss: -6.934479236602783\n",
      "Validation Loss: -6.943295955657959\n",
      "=== Epoch: 53 ===\n",
      "batch: 19, loss: -7.000945568084717\n",
      "batch: 39, loss: -5.499902248382568\n",
      "batch: 59, loss: -5.1334123611450195\n",
      "Validation Loss: -5.372155454423693\n",
      "=== Epoch: 54 ===\n",
      "batch: 19, loss: -7.139924049377441\n",
      "batch: 39, loss: -7.374932765960693\n",
      "batch: 59, loss: -7.535625457763672\n",
      "Validation Loss: -7.54930051167806\n",
      "=== Epoch: 55 ===\n",
      "batch: 19, loss: -7.7048659324646\n",
      "batch: 39, loss: -7.89346981048584\n",
      "batch: 59, loss: -8.122872352600098\n",
      "Validation Loss: -8.137826601664225\n",
      "=== Epoch: 56 ===\n",
      "batch: 19, loss: -8.322025299072266\n",
      "batch: 39, loss: -8.425168991088867\n",
      "batch: 59, loss: -8.247784614562988\n",
      "Validation Loss: -8.252712143792046\n",
      "=== Epoch: 57 ===\n",
      "batch: 19, loss: -8.212776184082031\n",
      "batch: 39, loss: -8.178881645202637\n",
      "batch: 59, loss: -8.38494873046875\n",
      "Validation Loss: -8.400687111748589\n",
      "=== Epoch: 58 ===\n",
      "batch: 19, loss: -8.235641479492188\n",
      "batch: 39, loss: -8.267157554626465\n",
      "batch: 59, loss: -7.9301557540893555\n",
      "Validation Loss: -8.32399876912435\n",
      "=== Epoch: 59 ===\n",
      "batch: 19, loss: -4.362585544586182\n",
      "batch: 39, loss: -4.027352333068848\n",
      "batch: 59, loss: -4.323105812072754\n",
      "Validation Loss: -4.349193096160889\n",
      "=== Epoch: 60 ===\n",
      "batch: 19, loss: -5.324522018432617\n",
      "batch: 39, loss: -5.787588119506836\n",
      "batch: 59, loss: -6.918570041656494\n",
      "Validation Loss: -6.891145494249132\n",
      "=== Epoch: 61 ===\n",
      "batch: 19, loss: -7.015007972717285\n",
      "batch: 39, loss: -7.177613735198975\n",
      "batch: 59, loss: -7.314838409423828\n",
      "Validation Loss: -7.3513970375061035\n",
      "=== Epoch: 62 ===\n",
      "batch: 19, loss: -7.478388786315918\n",
      "batch: 39, loss: -7.675145149230957\n",
      "batch: 59, loss: -7.898539066314697\n",
      "Validation Loss: -7.913917223612468\n",
      "=== Epoch: 63 ===\n",
      "batch: 19, loss: -8.10762882232666\n",
      "batch: 39, loss: -8.360712051391602\n",
      "batch: 59, loss: -7.527843475341797\n",
      "Validation Loss: -7.319175826178657\n",
      "=== Epoch: 64 ===\n",
      "batch: 19, loss: -5.881984710693359\n",
      "batch: 39, loss: -4.091148853302002\n",
      "batch: 59, loss: -4.08942174911499\n",
      "Validation Loss: -4.100026819441053\n",
      "=== Epoch: 65 ===\n",
      "batch: 19, loss: -4.420716285705566\n",
      "batch: 39, loss: -5.4454240798950195\n",
      "batch: 59, loss: -7.2647809982299805\n",
      "Validation Loss: -7.386932161119249\n"
     ]
    }
   ],
   "source": [
    "early_stop = 5\n",
    "min_loss = 10\n",
    "times = 0\n",
    "for i in range(epochs):\n",
    "    print(f\"=== Epoch: {i} ===\")\n",
    "    train_loop(train_dl, factor_VAE, optimizer)\n",
    "    val_loss = val_loop(val_dl, factor_VAE)\n",
    "    if val_loss < min_loss:\n",
    "        min_loss = val_loss\n",
    "        times = 0\n",
    "        best_model = factor_VAE.state_dict()\n",
    "        torch.save(best_model, \"best_model.pt\")\n",
    "    else:\n",
    "        times += 1\n",
    "    if times == early_stop:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation Loss: -8.395851254463196\n"
     ]
    }
   ],
   "source": [
    "factor_VAE = FactorVAE(\n",
    "    characteristic_size=characteristic_size,\n",
    "    stock_size=stock_size,\n",
    "    latent_size=latent_size,\n",
    "    factor_size=factor_size,\n",
    "    time_span=time_span,\n",
    "    gru_input_size=gru_input_size,\n",
    "    hidden_size=hidden_size\n",
    ").to(device)\n",
    "best_model_state = torch.load(\"best_model.pt\")\n",
    "factor_VAE.load_state_dict(best_model_state)\n",
    "factor_VAE.eval()\n",
    "loss = val_loop(test_dl, factor_VAE)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:28.162593Z",
     "start_time": "2024-03-06T12:03:22.117035Z"
    }
   },
   "execution_count": 23
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:34.246798Z",
     "start_time": "2024-03-06T12:03:28.164765Z"
    }
   },
   "outputs": [],
   "source": [
    "result = np.empty((len(test_features_dataset), stock_size))\n",
    "for i in range(len(test_features_dataset)):\n",
    "    tmp = factor_VAE.prediction(test_features_dataset[i].unsqueeze(0).to(device))\n",
    "    result[i] = tmp[0].squeeze().cpu().numpy()\n",
    "    '''\n",
    "    from torch.distributions import Normal\n",
    "    n = Normal(tmp[1], tmp[2])\n",
    "    n.sample()\n",
    "    '''"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(249, 487)"
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:34.278581Z",
     "start_time": "2024-03-06T12:03:34.249800Z"
    }
   },
   "execution_count": 25
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "ret = test_returns_dataset.cpu().numpy()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:34.293751Z",
     "start_time": "2024-03-06T12:03:34.280729Z"
    }
   },
   "execution_count": 26
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:34.325197Z",
     "start_time": "2024-03-06T12:03:34.295992Z"
    }
   },
   "outputs": [],
   "source": [
    "def calculate_returns(result, k, n):\n",
    "    num_days, num_stocks = result.shape\n",
    "    prev_selection = np.arange(num_stocks)  # 初始选择所有股票\n",
    "    returns_data = []\n",
    "    for i in range(num_days - 1):\n",
    "        today_returns = result[i]\n",
    "        sorted_indices = np.argsort(today_returns)[::-1]  # 按照收益率从高到低排序\n",
    "        # 选择前k个股票\n",
    "        today_selection = sorted_indices[:k]\n",
    "        # 计算前一天的选择股票与当天选择股票的交集\n",
    "        intersection_size = len(set(prev_selection) & set(today_selection))\n",
    "        # 如果交集大小小于 k - n，则交换掉前一天选择的股票中排名靠后的 n 支股票\n",
    "        if intersection_size < k - n:\n",
    "            prev_selection = today_selection\n",
    "        # 如果交集大小大于等于 k - n，则不进行交换\n",
    "        else:\n",
    "            today_selection = prev_selection\n",
    "        # 计算每天持有的 k 支股票的总收益率\n",
    "        daily_return = np.mean(ret[i+1][today_selection])  # 注意注意 ：这里是 ret 不是 result\n",
    "        returns_data.append(daily_return)\n",
    "    df = pd.DataFrame({\"daily_return\": returns_data})\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "     daily_return\n0        0.015230\n1       -0.010417\n2        0.023965\n3        0.000243\n4        0.007307\n..            ...\n243      0.012125\n244      0.003419\n245      0.005955\n246      0.001124\n247      0.001298\n\n[248 rows x 1 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>daily_return</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>0.015230</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>-0.010417</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>0.023965</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>0.000243</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>0.007307</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>243</th>\n      <td>0.012125</td>\n    </tr>\n    <tr>\n      <th>244</th>\n      <td>0.003419</td>\n    </tr>\n    <tr>\n      <th>245</th>\n      <td>0.005955</td>\n    </tr>\n    <tr>\n      <th>246</th>\n      <td>0.001124</td>\n    </tr>\n    <tr>\n      <th>247</th>\n      <td>0.001298</td>\n    </tr>\n  </tbody>\n</table>\n<p>248 rows × 1 columns</p>\n</div>"
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k = 400\n",
    "n = k // 10\n",
    "pred = calculate_returns(result, k, n)\n",
    "pred"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:04:15.579420Z",
     "start_time": "2024-03-06T12:04:15.537363Z"
    }
   },
   "execution_count": 34
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "     daily_return\n0        1.015230\n1        1.004654\n2        1.028731\n3        1.028981\n4        1.036500\n..            ...\n243      1.143701\n244      1.147611\n245      1.154446\n246      1.155743\n247      1.157244\n\n[248 rows x 1 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>daily_return</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>1.015230</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>1.004654</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>1.028731</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>1.028981</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>1.036500</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>243</th>\n      <td>1.143701</td>\n    </tr>\n    <tr>\n      <th>244</th>\n      <td>1.147611</td>\n    </tr>\n    <tr>\n      <th>245</th>\n      <td>1.154446</td>\n    </tr>\n    <tr>\n      <th>246</th>\n      <td>1.155743</td>\n    </tr>\n    <tr>\n      <th>247</th>\n      <td>1.157244</td>\n    </tr>\n  </tbody>\n</table>\n<p>248 rows × 1 columns</p>\n</div>"
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ARR = (1 + pred).cumprod()\n",
    "ARR"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:04:16.741817Z",
     "start_time": "2024-03-06T12:04:16.724827Z"
    }
   },
   "execution_count": 35
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "ARR.to_csv(f'{stock}.csv', index=None)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:03:34.542600Z",
     "start_time": "2024-03-06T12:03:34.498232Z"
    }
   },
   "execution_count": 30
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "portfolio_df_performance = pred"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:04:22.572721Z",
     "start_time": "2024-03-06T12:04:22.559212Z"
    }
   },
   "execution_count": 36
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\胡逸凡\\AppData\\Local\\Temp\\ipykernel_20996\\980728215.py:12: FutureWarning: 'm' is deprecated and will be removed in a future version, please use 'ME' instead.\n",
      "  monthly_statistics_df = alpha_df_performance[net_value_columns].resample('m').last()\n",
      "C:\\Users\\胡逸凡\\AppData\\Local\\Temp\\ipykernel_20996\\980728215.py:20: FutureWarning: 'y' is deprecated and will be removed in a future version, please use 'YE' instead.\n",
      "  yearly_statistics_df = alpha_df_performance[net_value_columns].resample('y').last()\n",
      "C:\\Users\\胡逸凡\\AppData\\Roaming\\Python\\Python310\\site-packages\\numpy\\core\\fromnumeric.py:3643: FutureWarning: The behavior of DataFrame.std with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)\n",
      "  return std(axis=axis, dtype=dtype, out=out, ddof=ddof, **kwargs)\n",
      "C:\\Users\\胡逸凡\\AppData\\Roaming\\Python\\Python310\\site-packages\\numpy\\core\\fromnumeric.py:3643: FutureWarning: The behavior of DataFrame.std with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)\n",
      "  return std(axis=axis, dtype=dtype, out=out, ddof=ddof, **kwargs)\n",
      "C:\\Users\\胡逸凡\\AppData\\Roaming\\Python\\Python310\\site-packages\\numpy\\core\\fromnumeric.py:3643: FutureWarning: The behavior of DataFrame.std with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)\n",
      "  return std(axis=axis, dtype=dtype, out=out, ddof=ddof, **kwargs)\n"
     ]
    }
   ],
   "source": [
    "alpha_df_performance = pd.DataFrame()\n",
    "alpha_df_performance['portfolio_daily_return'] = portfolio_df_performance['daily_return']\n",
    "alpha_df_performance['portfolio_net_value'] = (alpha_df_performance['portfolio_daily_return'] + 1).cumprod()\n",
    "\n",
    "net_value_columns = ['portfolio_net_value']\n",
    "\n",
    "alpha_statistics_df = pd.DataFrame(index=alpha_df_performance[net_value_columns].columns,\n",
    "                                    columns=[\"年化收益\", \"年化波动率\", \"最大回撤率\", \"夏普率\", \"Calmar\", \"IR\"])\n",
    "\n",
    "# alpha_df_performance.set_index(\"dt\", inplace=True)\n",
    "alpha_df_performance.index = pd.to_datetime(alpha_df_performance.index)\n",
    "monthly_statistics_df = alpha_df_performance[net_value_columns].resample('m').last()\n",
    "monthly_statistics_df = pd.concat([alpha_df_performance[:1][\n",
    "                                        ['portfolio_net_value']],\n",
    "                                    monthly_statistics_df])\n",
    "monthly_statistics_df = monthly_statistics_df.pct_change()\n",
    "monthly_statistics_df = monthly_statistics_df.dropna()\n",
    "monthly_statistics_df.index = monthly_statistics_df.index.date\n",
    "## TODO 补充第一年的数据\n",
    "yearly_statistics_df = alpha_df_performance[net_value_columns].resample('y').last()\n",
    "yearly_statistics_df = pd.concat([alpha_df_performance[:1][\n",
    "                                        ['portfolio_net_value']],\n",
    "                                    yearly_statistics_df])\n",
    "yearly_statistics_df = yearly_statistics_df.pct_change()\n",
    "yearly_statistics_df = yearly_statistics_df.dropna()\n",
    "yearly_statistics_df.index = yearly_statistics_df.index.date\n",
    "\n",
    "alpha_statistics_df.loc[:, \"年化收益\"] = np.mean(\n",
    "    (alpha_df_performance[net_value_columns].tail(1)) ** (252 / len(alpha_df_performance)) - 1)\n",
    "alpha_statistics_df.loc[:, \"年化波动率\"] = np.std(\n",
    "    alpha_df_performance[net_value_columns] / alpha_df_performance[net_value_columns].shift(1) - 1) * np.sqrt(\n",
    "    252)\n",
    "alpha_statistics_df.loc[:, \"累积收益\"] = np.mean(alpha_df_performance[net_value_columns].tail(1) - 1)\n",
    "alpha_statistics_df.loc[:, \"累积波动率\"] = np.std(\n",
    "    alpha_df_performance[net_value_columns] / alpha_df_performance[net_value_columns].shift(1) - 1)\n",
    "alpha_statistics_df.loc[:, \"最大回撤率\"] = np.min(\n",
    "    (alpha_df_performance[net_value_columns] - alpha_df_performance[net_value_columns].cummax()) /\n",
    "    alpha_df_performance[net_value_columns].cummax())\n",
    "alpha_statistics_df.loc[:, \"夏普率\"] = alpha_statistics_df[\"年化收益\"] / alpha_statistics_df[\"年化波动率\"]\n",
    "alpha_statistics_df.loc[:, \"Calmar\"] = alpha_statistics_df[\"年化收益\"] / abs(alpha_statistics_df[\"最大回撤率\"])\n",
    "alpha_statistics_df.loc[:, \"IR\"] = np.mean(\n",
    "    alpha_df_performance[net_value_columns] / alpha_df_performance[net_value_columns].shift(1) - 1) * np.sqrt(\n",
    "    252) / np.std(alpha_df_performance[net_value_columns] / alpha_df_performance[net_value_columns].shift(1) - 1)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:04:23.126999Z",
     "start_time": "2024-03-06T12:04:23.082989Z"
    }
   },
   "execution_count": 37
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "                         年化收益     年化波动率     最大回撤率       夏普率    Calmar  \\\nportfolio_net_value  0.159973  0.141798 -0.132042  1.128173  1.211524   \n\n                          IR      累积收益     累积波动率  \nportfolio_net_value  1.01311  0.157244  0.008932  ",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>年化收益</th>\n      <th>年化波动率</th>\n      <th>最大回撤率</th>\n      <th>夏普率</th>\n      <th>Calmar</th>\n      <th>IR</th>\n      <th>累积收益</th>\n      <th>累积波动率</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>portfolio_net_value</th>\n      <td>0.159973</td>\n      <td>0.141798</td>\n      <td>-0.132042</td>\n      <td>1.128173</td>\n      <td>1.211524</td>\n      <td>1.01311</td>\n      <td>0.157244</td>\n      <td>0.008932</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha_statistics_df"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T12:04:23.933915Z",
     "start_time": "2024-03-06T12:04:23.918712Z"
    }
   },
   "execution_count": 38
  },
  {
   "cell_type": "markdown",
   "source": [
    "FactorVAE\n",
    "ASR = ARR / AVol;   CR = ARR / abs(MDD)\n",
    "        ARR    AVol    MDD    ASR    CR    IR\n",
    "hs300 -0.048  0.134  -0.175 -0.355 -0.271 -0.348\n",
    "\n",
    "zz500  0.006  0.127  -0.147  0.047  0.041  0.112\n",
    "\n",
    "sp500  0.160  0.142  -0.132  1.128  1.211  1.013\n",
    "   \n",
    "nas100 0.356  0.159  -0.119  2.234  2.995  1.907"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.12 ('torch')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "deff6e117060b4680ad0a91c2cc9b11e790502a6f792302898866baebf8767a2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
